#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：LabelmeToYolo.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/1/7 上午8:56 
"""

import json
import os
import shutil
from typing import List, Tuple
import random


class LabelmeToYOLO:
	def __init__(self, input_dir: str, output_dir: str, train_ratio: float = 0.8):
		"""
		初始化Labelme到YOLO转换器。

		:param input_dir: 包含images和labels文件夹的输入目录路径。
		:param output_dir: 输出YOLO格式数据的目录路径。
		:param train_ratio: 训练集和测试集的比例，默认为0.8。
		"""
		self.input_dir = input_dir
		self.output_dir = output_dir
		self.train_ratio = train_ratio
		self.images_dir = os.path.join(input_dir, 'image')
		self.labels_dir = os.path.join(input_dir, 'label')
		self.train_images_dir = os.path.join(output_dir, 'train', 'image')
		self.train_labels_dir = os.path.join(output_dir, 'train', 'label')
		self.test_images_dir = os.path.join(output_dir, 'test', 'image')
		self.test_labels_dir = os.path.join(output_dir, 'test', 'label')
		os.makedirs(self.train_images_dir, exist_ok=True)
		os.makedirs(self.train_labels_dir, exist_ok=True)
		os.makedirs(self.test_images_dir, exist_ok=True)
		os.makedirs(self.test_labels_dir, exist_ok=True)

		self.classes = self._get_classes()

	def _get_classes(self) -> List[str]:
		"""
		从JSON文件中提取所有类别。

		:return: 类别列表。
		"""
		classes = set()
		for filename in os.listdir(self.labels_dir):
			if filename.endswith('.json'):
				with open(os.path.join(self.labels_dir, filename), 'r') as f:
					data = json.load(f)
					for shape in data['shapes']:
						classes.add(shape['label'])
		return sorted(list(classes))

	def _write_classes_file(self):
		"""
		将类别写入classes.txt文件。
		"""
		with open(os.path.join(self.output_dir, 'classes.txt'), 'w') as f:
			for cls in self.classes:
				f.write(cls + '\n')

	def _convert_shape(self, shape: dict, img_width: int, img_height: int) -> Tuple[float, float, float, float]:
		"""
		将Labelme的矩形形状转换为YOLO格式。

		:param shape: Labelme形状字典。
		:param img_width: 图像宽度。
		:param img_height: 图像高度。
		:return: YOLO格式的中心点坐标和宽高。
		"""
		x1, y1 = shape['points'][0]
		x2, y2 = shape['points'][1]
		x_center = (x1 + x2) / 2 / img_width
		y_center = (y1 + y2) / 2 / img_height
		width = (x2 - x1) / img_width
		height = (y2 - y1) / img_height
		return x_center, y_center, width, height

	def _convert_json_to_yolo(self, json_path: str, output_dir: str):
		"""
		将单个JSON文件转换为YOLO格式。

		:param json_path: JSON文件路径。
		:param output_dir: 输出目录。
		"""
		with open(json_path, 'r') as f:
			data = json.load(f)
			img_width, img_height = data['imageWidth'], data['imageHeight']
			txt_filename = os.path.splitext(os.path.basename(json_path))[0] + '.txt'
			with open(os.path.join(output_dir, txt_filename), 'w') as txt_file:
				for shape in data['shapes']:
					if shape['shape_type'] == 'rectangle':
						x_center, y_center, width, height = self._convert_shape(shape, img_width, img_height)
						class_id = self.classes.index(shape['label'])
						txt_file.write(f"{class_id} {x_center} {y_center} {width} {height}\n")

	def _copy_image(self, image_path: str, output_dir: str):
		"""
		复制图像文件到指定目录。

		:param image_path: 图像文件路径。
		:param output_dir: 输出目录。
		"""
		shutil.copy(image_path, output_dir)

	def convert_dataset(self):
		"""
		转换整个数据集为YOLO格式。
		"""
		self._write_classes_file()
		json_files = [f for f in os.listdir(self.labels_dir) if f.endswith('.json')]
		random.shuffle(json_files)
		train_size = int(len(json_files) * self.train_ratio)
		train_files = json_files[:train_size]
		test_files = json_files[train_size:]

		for json_file in train_files:
			json_path = os.path.join(self.labels_dir, json_file)
			self._convert_json_to_yolo(json_path, self.train_labels_dir)
			image_filename = os.path.splitext(json_file)[0] + '.jpg'
			image_path = os.path.join(self.images_dir, image_filename)
			if os.path.exists(image_path):
				self._copy_image(image_path, self.train_images_dir)

		for json_file in test_files:
			json_path = os.path.join(self.labels_dir, json_file)
			self._convert_json_to_yolo(json_path, self.test_labels_dir)
			image_filename = os.path.splitext(json_file)[0] + '.jpg'
			image_path = os.path.join(self.images_dir, image_filename)
			if os.path.exists(image_path):
				self._copy_image(image_path, self.test_images_dir)

		print(f"转换完成。训练集: {len(train_files)}个文件，测试集: {len(test_files)}个文件。")


# 使用示例
if __name__ == "__main__":
	input_dir = 'datasets/HELMET/source_data'
	output_dir = 'datasets/HELMET/yolo_data'
	converter = LabelmeToYOLO(input_dir, output_dir)
	converter.convert_dataset()
